package org.example;

import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyStatus;
import org.apache.rocketmq.client.producer.DefaultMQProducer;
import org.apache.rocketmq.client.producer.SendResult;
import org.apache.rocketmq.common.message.Message;
import org.apache.rocketmq.common.message.MessageExt;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;

public class FileProcess {
    private static final Map<String, Map<String, Integer>> receivedChunks = new HashMap<>();

    public void receiverFile(String path,MessageExt msg) {
        String key = msg.getKeys();
        String fileName = extractFileName(key);
        int chunkIndex = extractChunkIndex(key);
        int totalChunks = extractTotalChunks(key);

        // 根据文件名获取相应的哈希表
        Map<String, Integer> fileChunks = receivedChunks.computeIfAbsent(fileName, k -> new HashMap<>());

        if (!fileChunks.containsKey(key)) {
            saveChunkToFile(path, fileName, chunkIndex, totalChunks, msg.getBody());
            // 根据实际需求保存分片消息到哈希表
            fileChunks.put(key, 1);
            // 检查是否已接收到所有分片，如果是，则组装文件
            if (fileChunks.size() == totalChunks) {
                assembleFile(path, fileName, totalChunks);
                receivedChunks.remove(fileName);
            }
        }
    }


    public Boolean sendFile(DefaultMQProducer producer, String topic, String filePath) throws IOException {
        Path path = Paths.get(filePath);
        byte[] fileContent = Files.readAllBytes(path);
        int chunkSize = 1024 * 1024 * 4; // 1 MB per chunk
        String fileName = path.getFileName().toString();
        int totalChunks = (int) Math.ceil((double) fileContent.length / chunkSize);

        for (int i = 0; i < fileContent.length; i += chunkSize) {
            int endIndex = Math.min(i + chunkSize, fileContent.length);
            byte[] chunkData = new byte[endIndex - i];
            System.arraycopy(fileContent, i, chunkData, 0, chunkData.length);

            String key = String.format("FILE@@%s@@%d@@%d", fileName, i / chunkSize + 1, totalChunks);

            Message message = new Message(topic, "FILE", key, chunkData);
            try {
                SendResult sendResult = producer.send(message);
                System.out.printf("Send Result: %s%n", sendResult);
                Thread.sleep(500);
            } catch (Exception e) {
                e.printStackTrace();
                return false;
            }
        }
        return true;

    }

    private static String extractFileName(String key) {
        return key.split("@@")[1];
    }

    private static int extractChunkIndex(String key) {
        return Integer.parseInt(key.split("@@")[2]);
    }

    private static int extractTotalChunks(String key) {
        return Integer.parseInt(key.split("@@")[3]);
    }

    private static void saveChunkToFile(String path, String fileName, int chunkIndex, int totalChunks, byte[] chunkData) {
        // 根据实际需求将分片数据保存到文件
        String outputFile = String.format(path + "%s_%d_of_%d.txt", fileName, chunkIndex, totalChunks);
        try (FileOutputStream fos = new FileOutputStream(outputFile, true)) {
            fos.write(chunkData);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }


    private static void assembleFile(String path, String fileName, int totalChunks) {
        // 检查是否已接收到所有分片，如果是，则组装文件
        try (FileOutputStream fos = new FileOutputStream(path + fileName)) {
            for (int i = 1; i <= totalChunks; i++) {
                String chunkFile = String.format(path + "/%s_%d_of_%d.txt", fileName, i, totalChunks);
                byte[] chunkData = Files.readAllBytes(Paths.get(chunkFile));
                fos.write(chunkData);
                // 删除已组装的分片文件
                Files.deleteIfExists(Paths.get(chunkFile));
            }
            System.out.println("File"+"("+fileName+")"+" receives successfully!");
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}

